from flask import Flask, request, jsonify, send_file, Response
from word import GetWord
import os
import json
from werkzeug.utils import secure_filename as sf
from excel_processor import ExcelProcessor
import tempfile
from deepseek_api import DeepseekAPI

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 限制文件大小为16MB

# 确保上传目录存在
if not os.path.exists(app.config['UPLOAD_FOLDER']):
    os.makedirs(app.config['UPLOAD_FOLDER'])

@app.route('/')
def index():
    return app.send_static_file('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': '没有文件上传'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    filename = sf(file.filename)
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(filepath)

    try:
        # 判断文件类型
        if file.filename.endswith('.docx'):
            word_converter = GetWord()
            document = word_converter.get_document("", filepath)
            content = json.loads(document.page_content)
            
            # 删除临时文件
            os.remove(filepath)
            
            return jsonify({
                'success': True,
                'content': content
            })
        elif file.filename.endswith('.xlsx'):
            # 这里仅保存文件，不进行处理
            return jsonify({
                'success': True,
                'message': '已上传Excel文件',
                'filename': filename
            })
        else:
            os.remove(filepath)  # 删除不支持的文件
            return jsonify({'error': '不支持的文件类型，仅支持.docx和.xlsx文件'}), 400
    except Exception as e:
        # 确保清理临时文件
        if os.path.exists(filepath):
            os.remove(filepath)
        return jsonify({'error': str(e)}), 500

@app.route('/process-excel', methods=['POST'])
def process_excel():
    """处理已上传的Excel文件，生成prompt"""
    data = request.get_json()
    if not data or 'filename' not in data:
        return jsonify({'error': '请提供Excel文件名'}), 400
    
    filename = data['filename']
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], sf(filename))
    
    if not os.path.exists(filepath):
        return jsonify({'error': '文件不存在，请重新上传'}), 404
    
    try:
        processor = ExcelProcessor()
        result = processor.process_excel(filepath)
        
        if not result['success']:
            return jsonify({'error': result['message']}), 400
            
        # 返回临时文件路径，用于下载
        return jsonify({
            'success': True,
            'message': result['message'],
            'download_path': f"/download-excel/{os.path.basename(result['file_path'])}"
        })
    except Exception as e:
        return jsonify({'error': f'处理Excel失败: {str(e)}'}), 500

@app.route('/download-excel/<filename>', methods=['GET'])
def download_excel(filename):
    """下载生成的Excel文件"""
    # 在临时目录中查找文件
    temp_dir = tempfile.gettempdir()
    filepath = None
    
    # 搜索临时目录中的文件
    for root, dirs, files in os.walk(temp_dir):
        if filename in files:
            filepath = os.path.join(root, filename)
            break
    
    if not filepath or not os.path.exists(filepath):
        return jsonify({'error': '文件不存在或已过期'}), 404
        
    try:
        return send_file(filepath, as_attachment=True, 
                          download_name=filename)
    except Exception as e:
        return jsonify({'error': f'下载文件失败: {str(e)}'}), 500

@app.route('/api/chat-stream', methods=['POST'])
def chat_stream():
    """处理流式聊天请求"""
    try:
        data = request.json
        if not data:
            return jsonify({'error': '请求参数无效'}), 400
            
        messages = data.get('messages', [])
        max_tokens = data.get('max_tokens', 2048)
        
        if not messages or not isinstance(messages, list) or len(messages) == 0:
            return jsonify({'error': '消息格式不正确'}), 400
        
        # 创建DeepSeek API客户端
        api = DeepseekAPI()
        
        def generate():
            # 获取流式响应
            stream_response = api.chat_completion(messages, max_tokens=max_tokens, stream=True)
            
            # 处理每个响应块
            for chunk in stream_response:
                if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
                    # 输出思维链内容
                    yield chunk.choices[0].delta.reasoning_content
                elif hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
                    # 输出最终回答内容
                    yield chunk.choices[0].delta.content
        
        # 返回流式响应
        return Response(generate(), mimetype='text/plain')
    except Exception as e:
        return jsonify({'error': f'处理聊天请求失败: {str(e)}'}), 500

@app.route('/review-document', methods=['POST'])
def review_document():
    """文档审核功能，使用DeepSeek R1模型"""
    try:
        if 'wordFile' not in request.files or 'excelFile' not in request.files:
            return jsonify({'error': '缺少必要的文件'}), 400
            
        word_file = request.files['wordFile']
        excel_file = request.files['excelFile']
        model_type = request.form.get('modelType', 'deepseek')
        
        # 保存上传的文件
        word_filename = sf(word_file.filename)
        excel_filename = sf(excel_file.filename)
        
        word_filepath = os.path.join(app.config['UPLOAD_FOLDER'], word_filename)
        excel_filepath = os.path.join(app.config['UPLOAD_FOLDER'], excel_filename)
        
        word_file.save(word_filepath)
        excel_file.save(excel_filepath)
        
        # 解析Word文档内容
        word_converter = GetWord()
        document = word_converter.get_document("", word_filepath)
        word_content = json.loads(document.page_content)
        
        # 处理Excel文件
        try:
            import pandas as pd
            df = pd.read_excel(excel_filepath)
            
            # 确保结果和thinking列存在
            if 'result' not in df.columns:
                df['result'] = ''
            if 'thinking' not in df.columns:
                df['thinking'] = ''
                
            # 创建DeepSeek API客户端
            api = DeepseekAPI()
            
            # 处理每一行
            for index, row in df.iterrows():
                # 构建提示词(从excel的prompt列获取)
                prompt = row.get('prompt', '')
                if not prompt:
                    continue
                    
                # 替换提示词中的占位符
                for key in word_content:
                    placeholder = f"{{{{document.{key}}}}}"
                    if placeholder in prompt:
                        prompt = prompt.replace(placeholder, word_content[key])
                
                # 构建消息
                messages = [{"role": "user", "content": prompt}]
                
                # 调用API
                response = api.chat_completion(messages)
                
                # 提取结果
                thinking = response.choices[0].message.reasoning_content
                result = response.choices[0].message.content
                
                # 更新Excel
                df.at[index, 'thinking'] = thinking
                df.at[index, 'result'] = result
            
            # 保存结果到临时文件
            import tempfile
            output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx')
            df.to_excel(output_file.name, index=False)
            
            # 清理原始上传文件
            os.remove(word_filepath)
            os.remove(excel_filepath)
            
            return jsonify({
                'success': True,
                'message': '文档审核完成',
                'download_path': f"/download-excel/{os.path.basename(output_file.name)}"
            })
                
        except Exception as e:
            return jsonify({'error': f'处理Excel失败: {str(e)}'}), 500
        
    except Exception as e:
        return jsonify({'error': f'审核文档失败: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)
